from PIL import Image
import shutil,pytesseract
import requests,os
from ua import User_Agent
from aip import AipOcr
import fitz
class To_str(object):
    def __init__(self,url,name):
        self.headers = {'User_Agent':User_Agent}
        self.request = requests
        self.baseurl = url
        self.name = name
        #百度OCR账号，密码，密钥
        self.config= {
        'appId': '',                    
        'apiKey': '',
        'secretKey': ''
    }
        self.client=AipOcr(**self.config)
        #选择一种识别方式，百度OCR为在线方式，线下识别为tessract库
        self.OCR_ONLINE = True                               
    def img_to_str_tesseract(self,image_path, lang='chi_sim'):
        print(pytesseract.image_to_string(Image.open(image_path), lang))
        return pytesseract.image_to_string(Image.open(image_path), lang)

    def write_file(self,path, text, ftype, debug=False):
        with open(path, ftype) as f:
            if debug:
                print("write", len(text))
            f.write(text.encode("gbk", 'ignore').decode("gbk", "ignore"))



    def img_to_str_baidu(self,image_path):
        with open(image_path, 'rb') as fp:
            image = fp.read()
            result = self.client.basicGeneral(image)
            if 'words_result' in result:
                return '\n'.join([w['words'] for w in result['words_result']])
        return ""

    def pdf_img(self):
        # print(url, name)
        path = os.path.abspath(os.path.dirname(__file__)).replace('\\', '/')
        filepath = path + '/' + 'pdf'
        if not os.path.isdir(filepath):
            os.mkdir(filepath)
        response = requests.get(self.baseurl)
        # print(response.status_code)
        response.encoding = response.apparent_encoding
        with open(filepath + '/{}.pdf'.format(self.name), 'wb+')as fp:
            fp.write(response.content)
            print('下载成功')
        doc = fitz.open(filepath + '/%s.pdf' % self.name)
        for pg in range(doc.pageCount):
            page = doc[pg]
            rotate = int(0)
            # 每个尺寸的缩放系数为2，这将为我们生成分辨率提高四倍的图像。
            zoom_x = 2.0
            zoom_y = 2.0
            trans = fitz.Matrix(zoom_x, zoom_y).preRotate(rotate)
            pm = page.getPixmap(matrix=trans, alpha=False)
            if not os.path.isdir(path + '/png'):
                os.mkdir(path + '/png')
            pm.writePNG(path + '/png' + '/%s.png' % pg)
        doc.close()
        list = os.listdir(path + '/png')
        outpath = path + '/txt'
        if not os.path.isdir(outpath):
            os.mkdir(outpath)
        for i in list:
            path_img = path + '/png' + '/{}'.format(i)
            if self.OCR_ONLINE:
                self.write_file(outpath + '/{}.txt'.format(self.name), self.img_to_str_baidu(path_img), 'a+')
            else:
                self.write_file(outpath + '/{}.txt'.format(self.name), self.img_to_str_tesseract(path_img), 'a+')
        with open(outpath + '/%s.txt' % self.name, 'r+')as fp:
            content = fp.readlines()
        str = ''
        for par in content:
            str = str + par
        strs = str.replace(' ', '').replace('\t', '').replace('\n', '').replace('\r', '')
        shutil.rmtree(filepath)
        shutil.rmtree(path + '/png')
        shutil.rmtree(path + '/txt')
        return strs